#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import asyncio
import threading
import time
from flask import Flask, render_template, request
from flask_socketio import SocketIO, emit
from flask_cors import CORS
import zmq
import zmq.asyncio

app = Flask(__name__)
app.config['SECRET_KEY'] = 'webrtc_demo_secret'
CORS(app, cors_allowed_origins="*")
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

# 存储连接的客户端
clients = {}
remote_servers = {}

# ZMQ配置
ZMQ_PORT = 5555
zmq_context = None
zmq_socket = None

def init_zmq(port=None):
    """初始化ZMQ服务器"""
    global zmq_context, zmq_socket, ZMQ_PORT
    if port:
        ZMQ_PORT = port
    
    zmq_context = zmq.Context()
    zmq_socket = zmq_context.socket(zmq.ROUTER)
    zmq_socket.bind(f"tcp://*:{ZMQ_PORT}")
    print(f"ZMQ服务器启动在端口 {ZMQ_PORT}")

def zmq_message_handler():
    """处理ZMQ消息的线程"""
    while True:
        try:
            # 接收ZMQ消息，DEALER发送的消息格式为 [message]
            parts = zmq_socket.recv_multipart(zmq.NOBLOCK)
            
            print(f"收到ZMQ消息部分数: {len(parts)}")
            
            if len(parts) == 1:
                # DEALER直接发送消息: [message]
                message = parts[0]
                identity = b'unknown'
            elif len(parts) == 2:
                # DEALER格式: [identity, message]
                identity, message = parts
            elif len(parts) == 3:
                # ROUTER格式: [identity, empty, message]
                identity, empty, message = parts
            else:
                print(f"未知ZMQ消息格式: {len(parts)}个部分")
                continue
                
            try:
                message_str = message.decode('utf-8').strip()
                if not message_str:
                    print("收到空ZMQ消息，跳过")
                    continue
                    
                data = json.loads(message_str)
                print(f"收到ZMQ消息: {data}")
            except json.JSONDecodeError as e:
                print(f"JSON解析错误: {e}, 消息内容: {repr(message)}")
                continue
            except UnicodeDecodeError as e:
                print(f"消息解码错误: {e}, 消息内容: {repr(message)}")
                continue
            
            # 处理不同类型的消息
            if data['type'] == 'register':
                server_id = data['server_id']
                remote_servers[identity] = {
                    'id': server_id,
                    'identity': identity,
                    'last_seen': time.time()
                }
                print(f"远程服务器注册成功: {server_id} (identity: {identity})")
                print(f"当前注册的服务器数量: {len(remote_servers)}")
                
            elif data['type'] == 'heartbeat':
                server_id = data['server_id']
                # 更新心跳时间，如果服务器未注册则自动注册
                if identity not in remote_servers:
                    remote_servers[identity] = {
                        'id': server_id,
                        'identity': identity,
                        'last_seen': time.time()
                    }
                    print(f"通过心跳自动注册服务器: {server_id}")
                else:
                    remote_servers[identity]['last_seen'] = time.time()
                
            elif data['type'] == 'offer' or data['type'] == 'answer' or data['type'] == 'ice-candidate':
                # 转发WebRTC信令消息到前端
                target_client = data.get('target')
                if target_client in clients:
                    socketio.emit('webrtc_signal', data, room=clients[target_client]['sid'])
                    print(f"转发信令消息到客户端: {target_client}")
                    
        except zmq.Again:
            time.sleep(0.01)
        except Exception as e:
            print(f"ZMQ消息处理错误: {e}")
            time.sleep(0.1)

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@socketio.on('connect')
def handle_connect():
    """客户端连接"""
    client_id = request.sid
    clients[client_id] = {
        'sid': client_id,
        'connected_time': time.time()
    }
    print(f"客户端连接: {client_id}")
    emit('connected', {'client_id': client_id})

@socketio.on('disconnect')
def handle_disconnect():
    """客户端断开连接"""
    client_id = request.sid
    if client_id in clients:
        del clients[client_id]
    print(f"客户端断开连接: {client_id}")

@socketio.on('webrtc_signal')
def handle_webrtc_signal(data):
    """处理WebRTC信令消息"""
    client_id = request.sid
    data['source'] = client_id
    
    print(f"收到WebRTC信令: {data['type']} from {client_id}")
    
    # 转发到远程服务器
    if remote_servers:
        # 清理过期的服务器
        current_time = time.time()
        expired_servers = []
        for identity, server in remote_servers.items():
            if current_time - server['last_seen'] > 60:  # 60秒超时
                expired_servers.append(identity)
        
        for identity in expired_servers:
            print(f"清理过期服务器: {remote_servers[identity]['id']}")
            del remote_servers[identity]
        
        # 选择最近活跃的远程服务器
        if remote_servers:
            # 按最后活跃时间排序，选择最新的
            active_servers = sorted(remote_servers.items(), 
                                  key=lambda x: x[1]['last_seen'], reverse=True)
            server_identity = active_servers[0][0]
            server_id = active_servers[0][1]['id']
            
            message = json.dumps(data)
            zmq_socket.send_multipart([server_identity, message.encode('utf-8')])
            print(f"转发信令消息到远程服务器 {server_id}: {data['type']}")
        else:
            print("没有可用的远程服务器")
            emit('error', {'message': '没有可用的远程服务器'})
    else:
        print("没有可用的远程服务器")
        emit('error', {'message': '没有可用的远程服务器'})

@socketio.on('get_servers')
def handle_get_servers():
    """获取可用的远程服务器列表"""
    server_list = []
    current_time = time.time()
    
    # 清理超时的服务器
    expired_servers = []
    for identity, server in remote_servers.items():
        if current_time - server['last_seen'] > 30:  # 30秒超时
            expired_servers.append(identity)
        else:
            server_list.append({
                'id': server['id'],
                'status': 'online'
            })
    
    for identity in expired_servers:
        del remote_servers[identity]
    
    emit('servers_list', {'servers': server_list})

def start_zmq_handler():
    """启动ZMQ消息处理线程"""
    zmq_thread = threading.Thread(target=zmq_message_handler, daemon=True)
    zmq_thread.start()

if __name__ == '__main__':
    # 初始化ZMQ
    init_zmq()
    
    # 启动ZMQ消息处理线程
    start_zmq_handler()
    
    print("信令服务器启动...")
    print("WebSocket服务: http://localhost:5000")
    print(f"ZMQ服务: tcp://localhost:{ZMQ_PORT}")
    
    # 启动Flask-SocketIO服务器
    socketio.run(app, host='0.0.0.0', port=5000, debug=False)